{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "1WNN7lAwNPRD",
      "metadata": {
        "id": "1WNN7lAwNPRD"
      },
      "source": [
        "# Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "OqBDNX7tnO7w",
      "metadata": {
        "id": "OqBDNX7tnO7w"
      },
      "outputs": [],
      "source": [
        "! pip install --upgrade --quiet langchain langchain_google_community langchain_google_vertexai google-cloud-discoveryengine google_cloud_aiplatform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fsG9CSrOFzHt",
      "metadata": {
        "id": "fsG9CSrOFzHt"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "O1yHvoyXFaPc",
      "metadata": {
        "id": "O1yHvoyXFaPc"
      },
      "outputs": [],
      "source": [
        "# Authenticate with Colab (if you're using Google Colab)\n",
        "\n",
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()\n",
        "\n",
        "from google.auth import default\n",
        "creds, _ = default()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "NbO6qHY2IKDg",
      "metadata": {
        "id": "NbO6qHY2IKDg"
      },
      "outputs": [],
      "source": [
        "# Set your project and region\n",
        "PROJECT_ID = \"my-project\"  # @param {type:\"string\"}\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "LOCATION = \"global\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ruFJiEYpGQ50",
      "metadata": {
        "id": "ruFJiEYpGQ50"
      },
      "outputs": [],
      "source": [
        "# This cell is only needed if you run Google Colab\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=REGION)\n",
        "! gcloud config set project {PROJECT_ID}\n",
        "! gcloud auth application-default login\n",
        "! gcloud auth application-default set-quota-project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "DH-Hsd-pNSN4",
      "metadata": {
        "id": "DH-Hsd-pNSN4"
      },
      "source": [
        "# Hallucinations"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "A7lPJMqOM6Fn",
      "metadata": {
        "id": "A7lPJMqOM6Fn"
      },
      "source": [
        "## How prompt engineering helps to avoid hallucinations"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "DEf4Rxl6F_Rw",
      "metadata": {
        "id": "DEf4Rxl6F_Rw"
      },
      "source": [
        "Hallucinations are a real problem for LLMs. Since the LLM completes text that is plausible, it can be tricked into creating plausible, but completely false responses. Please take note that we are intentionally using an \"old\" model here, as newer model have much better built-in guardrailes. Here's a simple example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fon_e0FgF-hA",
      "metadata": {
        "id": "fon_e0FgF-hA"
      },
      "outputs": [],
      "source": [
        "from langchain.chains import LLMChain\n",
        "from langchain_core.prompts.prompt import PromptTemplate\n",
        "from langchain_google_vertexai import VertexAI\n",
        "\n",
        "\n",
        "llm = VertexAI(model_name=\"gemini-1.5-pro-001\", temperature=0.8, max_output_tokens=128)\n",
        "\n",
        "template = \"\"\"Describe {plant}.\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "prompt_template = PromptTemplate(\n",
        "   input_variables=[\"plant\"],\n",
        "   template=template,\n",
        ")\n",
        "\n",
        "\n",
        "chain = LLMChain(llm=llm, prompt=prompt_template)\n",
        "chain.run(plant=\"black cucumbers\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f8JLGjw4M8tI",
      "metadata": {
        "id": "f8JLGjw4M8tI"
      },
      "source": [
        "Let's look at a simple prompt adjustment that helps to avoid hallucinations:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "Az5cbdYNGCOMpfjIjTBBjN1u",
      "metadata": {
        "id": "Az5cbdYNGCOMpfjIjTBBjN1u",
        "tags": []
      },
      "outputs": [],
      "source": [
        "from langchain.chains import LLMChain\n",
        "from langchain_core.prompts.prompt import PromptTemplate\n",
        "from langchain_google_vertexai import VertexAI\n",
        "\n",
        "\n",
        "llm = VertexAI(model_name=\"gemini-1.0-pro-001\", temperature=0.8, max_output_tokens=128)\n",
        "\n",
        "template = \"\"\"Describe {plant}.\n",
        "\n",
        "\n",
        "First, think whether {plant} exist.\n",
        "If {plant} doesn't exist, answer \"I don't have enough information about {plant}\".\n",
        "Otherwise, give its title, a short summary and then talk about origin and cultivation.\n",
        "After that, describe its physical characteristics.\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "prompt_template = PromptTemplate(\n",
        "   input_variables=[\"plant\"],\n",
        "   template=template,\n",
        ")\n",
        "\n",
        "\n",
        "chain = LLMChain(llm=llm, prompt=prompt_template)\n",
        "chain.run(plant=\"black cucumbers\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "h87kNLX_ULyS",
      "metadata": {
        "id": "h87kNLX_ULyS"
      },
      "source": [
        "# Vertex AI Agent Builder"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "aVw2HajqUNLg",
      "metadata": {
        "id": "aVw2HajqUNLg"
      },
      "source": [
        "## Using Vertex AI Agent Builder as a tool"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "PgAgN1CzUOyJ",
      "metadata": {
        "id": "PgAgN1CzUOyJ"
      },
      "source": [
        "Vertex AI Agent Builder provides an end-to-end experience for build out-of-the-box RAG agents. You will find relevant documents and be able to prepare a final answer with generative AI. You can use it as a tool. We will look in Chapter 9 what tools are and how they help to build powerful generative AI agents, but all we need to know for now is that Vertex AI Agent Builder is an interface that takes a question as an input and returns an answer.\n",
        "\n",
        "Check the [Google Cloud Documentation](https://cloud.google.com/products/agent-builder?e=0&hl=en#use-vertex-ai-search-for-out-of-the-box-rag-for-your-agents-and-apps) for a deep discussion of Vertex AI Agent Builder for Search."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "P5QaB5BfUxyX",
      "metadata": {
        "id": "P5QaB5BfUxyX"
      },
      "source": [
        "Vertex AI Agent Builder uses so-called Data Stores for running its code. Follow the instructions in the book on how to set up your datastore, and after you created it, you can use it with LangChain:"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "pFHW3Vt1u_gA",
      "metadata": {
        "id": "pFHW3Vt1u_gA"
      },
      "source": [
        "### Creating the Datastore"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fNlOt8vSvB1q",
      "metadata": {
        "id": "fNlOt8vSvB1q"
      },
      "source": [
        "First, we need to create the datastore. This is the place, where you will hold the data that flows into your RAG. Confusingly, it uses something called a Discoveryengine. Check the documentation for details on datastores [here](https://cloud.google.com/generative-ai-app-builder/docs/create-data-store-es)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "zsqfVLRDvBdP",
      "metadata": {
        "id": "zsqfVLRDvBdP"
      },
      "outputs": [],
      "source": [
        "# Let's start with with some helper functions\n",
        "def create_data_store(\n",
        "    project_id: str,\n",
        "    location: str,\n",
        "    data_store_id: str,\n",
        ") -> str:\n",
        "    #  For more information, refer to:\n",
        "    # https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store\n",
        "    client_options = (\n",
        "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
        "        if location != \"global\"\n",
        "        else None\n",
        "    )\n",
        "\n",
        "    # Create a client\n",
        "    client = discoveryengine.DataStoreServiceClient(client_options=client_options)\n",
        "\n",
        "    # The full resource name of the collection\n",
        "    # e.g. projects/{project}/locations/{location}/collections/default_collection\n",
        "    parent = client.collection_path(\n",
        "        project=project_id,\n",
        "        location=location,\n",
        "        collection=\"default_collection\",\n",
        "    )\n",
        "\n",
        "    data_store = discoveryengine.DataStore(\n",
        "        display_name=\"My Data Store\",\n",
        "        # Options: GENERIC, MEDIA, HEALTHCARE_FHIR\n",
        "        industry_vertical=discoveryengine.IndustryVertical.GENERIC,\n",
        "        # Options: SOLUTION_TYPE_RECOMMENDATION, SOLUTION_TYPE_SEARCH, SOLUTION_TYPE_CHAT, SOLUTION_TYPE_GENERATIVE_CHAT\n",
        "        solution_types=[discoveryengine.SolutionType.SOLUTION_TYPE_SEARCH],\n",
        "        # TODO(developer): Update content_config based on data store type.\n",
        "        # Options: NO_CONTENT, CONTENT_REQUIRED, PUBLIC_WEBSITE\n",
        "        content_config=discoveryengine.DataStore.ContentConfig.CONTENT_REQUIRED,\n",
        "    )\n",
        "\n",
        "    request = discoveryengine.CreateDataStoreRequest(\n",
        "        parent=parent,\n",
        "        data_store_id=data_store_id,\n",
        "        data_store=data_store,\n",
        "        # Optional: For Advanced Site Search Only\n",
        "        # create_advanced_site_search=True,\n",
        "    )\n",
        "\n",
        "    # Make the request\n",
        "    operation = client.create_data_store(request=request)\n",
        "\n",
        "    print(f\"Waiting for operation to complete: {operation.operation.name}\")\n",
        "    response = operation.result()\n",
        "\n",
        "    # Once the operation is complete,\n",
        "    # get information from operation metadata\n",
        "    metadata = discoveryengine.CreateDataStoreMetadata(operation.metadata)\n",
        "\n",
        "    # Handle the response\n",
        "    print(response)\n",
        "    print(metadata)\n",
        "\n",
        "    return operation.operation.name\n",
        "\n",
        "def import_documents(\n",
        "    project_id: str,\n",
        "    location: str,\n",
        "    data_store_id: str,\n",
        "    gcs_uri: str,\n",
        "    ):\n",
        "    # Create a client\n",
        "    client_options = (\n",
        "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
        "        if location != \"global\"\n",
        "        else None\n",
        "    )\n",
        "    client = discoveryengine.DocumentServiceClient(client_options=client_options)\n",
        "\n",
        "    # The full resource name of the search engine branch.\n",
        "    # e.g. projects/{project}/locations/{location}/dataStores/{data_store_id}/branches/{branch}\n",
        "    parent = client.branch_path(\n",
        "        project=project_id,\n",
        "        location=location,\n",
        "        data_store=data_store_id,\n",
        "        branch=\"default_branch\",\n",
        "    )\n",
        "    source_documents = [f\"{gcs_uri}/*\"]\n",
        "\n",
        "    request = discoveryengine.ImportDocumentsRequest(\n",
        "        parent=parent,\n",
        "        gcs_source=discoveryengine.GcsSource(\n",
        "            input_uris=source_documents, data_schema=\"content\"\n",
        "        ),\n",
        "        # Options: `FULL`, `INCREMENTAL`\n",
        "        reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,\n",
        "    )\n",
        "\n",
        "    # Make the request\n",
        "    operation = client.import_documents(request=request)\n",
        "\n",
        "    response = operation.result()\n",
        "\n",
        "    # Once the operation is complete,\n",
        "    # get information from operation metadata\n",
        "    metadata = discoveryengine.ImportDocumentsMetadata(operation.metadata)\n",
        "\n",
        "    # Handle the response\n",
        "    return operation.operation.name\n",
        "\n",
        "def create_engine(\n",
        "    project_id: str, location: str, data_store_name: str, data_store_id: str\n",
        "  ):\n",
        "    # Create a client\n",
        "    client_options = (\n",
        "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
        "        if location != \"global\"\n",
        "        else None\n",
        "    )\n",
        "    client = discoveryengine.EngineServiceClient(client_options=client_options)\n",
        "\n",
        "    # Initialize request argument(s)\n",
        "    config = discoveryengine.Engine.SearchEngineConfig(\n",
        "        search_tier=\"SEARCH_TIER_ENTERPRISE\", search_add_ons=[\"SEARCH_ADD_ON_LLM\"]\n",
        "    )\n",
        "\n",
        "    engine = discoveryengine.Engine(\n",
        "        display_name=data_store_name,\n",
        "        solution_type=\"SOLUTION_TYPE_SEARCH\",\n",
        "        industry_vertical=\"GENERIC\",\n",
        "        data_store_ids=[data_store_id],\n",
        "        search_engine_config=config,\n",
        "    )\n",
        "\n",
        "    request = discoveryengine.CreateEngineRequest(\n",
        "        parent=discoveryengine.DataStoreServiceClient.collection_path(\n",
        "            project_id, location, \"default_collection\"\n",
        "        ),\n",
        "        engine=engine,\n",
        "        engine_id=engine.display_name,\n",
        "    )\n",
        "\n",
        "    # Make the request\n",
        "    operation = client.create_engine(request=request)\n",
        "    response = operation.result(timeout=90)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "sgBwubY9UTj2",
      "metadata": {
        "id": "sgBwubY9UTj2"
      },
      "outputs": [],
      "source": [
        "from google.cloud import discoveryengine_v1alpha as discoveryengine\n",
        "from google.api_core.client_options import ClientOptions\n",
        "\n",
        "# The datastore name can only contain lowercase letters, numbers, and hyphens\n",
        "DATASTORE_NAME = \"movie-database-maxtsc\"\n",
        "DATASTORE_ID = f\"{DATASTORE_NAME}-id\"\n",
        "COLLECTION = create_data_store(PROJECT_ID, LOCATION, DATASTORE_ID)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "xFdWoDUbBykD",
      "metadata": {
        "id": "xFdWoDUbBykD"
      },
      "outputs": [],
      "source": [
        "# Now we need to also set up a search engine on top of the datastore so we can use enterprise features\n",
        "create_engine(PROJECT_ID, LOCATION, DATASTORE_NAME, DATASTORE_ID)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "XKF9qg0L_ZWK",
      "metadata": {
        "id": "XKF9qg0L_ZWK"
      },
      "outputs": [],
      "source": [
        "# Now we import documents. We load a folder that contains arxiv articles. This will take some time ☕🥪\n",
        "import_documents(PROJECT_ID, LOCATION, DATASTORE_ID,\"gs://cloud-samples-data/gen-app-builder/search/arxiv\")\n",
        "\n",
        "# Don't worry if it times out, you can check the status in the Agent Builder console and use the tool once it's complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "EtEIn7LJBcZC",
      "metadata": {
        "id": "EtEIn7LJBcZC"
      },
      "outputs": [],
      "source": [
        "from langchain_google_community import VertexAISearchSummaryTool\n",
        "\n",
        "vertex_search = VertexAISearchSummaryTool(\n",
        "  project_id=PROJECT_ID, location_id=LOCATION,\n",
        "  data_store_id=DATASTORE_ID,\n",
        "  name=\"Vertex AI Search\", description=\"\")\n",
        "\n",
        "query = \"what's the best way to improve the performance of LLM prompts?\"\n",
        "\n",
        "print(vertex_search.invoke(query))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "SzXL2GYk0MHk",
      "metadata": {
        "id": "SzXL2GYk0MHk"
      },
      "source": [
        "### LangChain Documents"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "hNqUBBQiH9xT",
      "metadata": {
        "id": "hNqUBBQiH9xT"
      },
      "source": [
        "Before we dive into building the RAG, it's important to understand how LangChain documents work. Below is a simple example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "BtAlAwUeNdGk",
      "metadata": {
        "id": "BtAlAwUeNdGk"
      },
      "outputs": [],
      "source": [
        "from langchain_core.documents import Document\n",
        "\n",
        "doc = Document(page_content=\"my page\",\n",
        "              metadata={\"source_id\": \"example.pdf\", \"page\": 1})\n",
        "print(doc.page_content)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9ewFgX0w0OpL",
      "metadata": {
        "id": "9ewFgX0w0OpL"
      },
      "source": [
        "### Using Vertex AI Agent Builder Search as retriever"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "sWUaU-qBHmA2",
      "metadata": {
        "id": "sWUaU-qBHmA2"
      },
      "source": [
        "Now we can also use our newly built tool as a retriever to have more control over the context it returns."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "D-QQISqNVG0h",
      "metadata": {
        "id": "D-QQISqNVG0h"
      },
      "outputs": [],
      "source": [
        "from langchain_google_community import VertexAISearchRetriever\n",
        "\n",
        "vertex_search_retriever = VertexAISearchRetriever(\n",
        "   project_id=PROJECT_ID,\n",
        "   location_id=LOCATION,\n",
        "   data_store_id=DATASTORE_ID,\n",
        "   max_documents=3,\n",
        ")\n",
        "\n",
        "\n",
        "result = vertex_search_retriever.invoke(query)\n",
        "print(len(result))\n",
        "\n",
        "for doc in result:\n",
        "   print(len(doc.page_content), doc.metadata)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "_P6YeYg60TJi",
      "metadata": {
        "id": "_P6YeYg60TJi"
      },
      "source": [
        "## Building a RAG"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "TVsk0k1iHss5",
      "metadata": {
        "id": "TVsk0k1iHss5"
      },
      "source": [
        "Finally, with everything in place, we can build a beautiful RAG application."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "GdaH85bmVYnM",
      "metadata": {
        "id": "GdaH85bmVYnM"
      },
      "outputs": [],
      "source": [
        "from langchain_core.output_parsers import StrOutputParser\n",
        "from langchain_core.prompts import ChatPromptTemplate\n",
        "from langchain_core.runnables import RunnablePassthrough\n",
        "from langchain_google_vertexai import VertexAI\n",
        "\n",
        "\n",
        "template = \"\"\"Answer the question based only on the following context:\n",
        "{context}\n",
        "\n",
        "\n",
        "Question: {question}\n",
        "\"\"\"\n",
        "prompt = ChatPromptTemplate.from_template(template)\n",
        "\n",
        "\n",
        "llm = VertexAI(model_name=\"gemini-1.5-flash-001\", temperature=0.4, max_output_tokens=512)\n",
        "\n",
        "chain = (\n",
        "   {\"context\": vertex_search_retriever, \"question\": RunnablePassthrough()}\n",
        "   | prompt\n",
        "   | llm\n",
        "   | StrOutputParser()\n",
        ")\n",
        "\n",
        "\n",
        "chain.invoke(\"How do I make my LLM prompts perform better?\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5H8ZpZiTiqXB",
      "metadata": {
        "id": "5H8ZpZiTiqXB"
      },
      "source": [
        "## Query expansion"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "SUtpCA69X3XX",
      "metadata": {
        "id": "SUtpCA69X3XX"
      },
      "source": [
        "Query expansion is an attempt to reformulate queries to get a better (and broader) list of chunks.  One way of doing this is using `MultiQueryRetriever` component. Let's see how the amount of chunks increases after query expansion:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6arSZeUpXaul",
      "metadata": {
        "id": "6arSZeUpXaul"
      },
      "outputs": [],
      "source": [
        "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
        "\n",
        "\n",
        "retriever_with_expansion = MultiQueryRetriever.from_llm(\n",
        "   retriever=vertex_search_retriever, llm=llm\n",
        ")\n",
        "\n",
        "result = vertex_search_retriever.invoke(query)\n",
        "print(len(result))\n",
        "\n",
        "\n",
        "result_expansion = retriever_with_expansion.invoke(query)\n",
        "print(len(result_expansion))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "gyWTNb1WisPZ",
      "metadata": {
        "id": "gyWTNb1WisPZ"
      },
      "source": [
        "## Filtering out irrelevant chunks"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "CmS9vCNFiX9_",
      "metadata": {
        "id": "CmS9vCNFiX9_"
      },
      "source": [
        "Sometimes retrieval returns documents that are not really relevant, and they confuse the model. You can filter them by making additional pass through each document and asking the LLM to evaluate its relevance with `LLMChainFilter`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "gYeL6AkPhkQe",
      "metadata": {
        "id": "gYeL6AkPhkQe"
      },
      "outputs": [],
      "source": [
        "from langchain.retrievers.document_compressors import LLMChainFilter\n",
        "\n",
        "# Retrieve many documents from the retrieval\n",
        "vertex_search_retriever_many = VertexAISearchRetriever(\n",
        "   project_id=project_id,\n",
        "   location_id=location_id,\n",
        "   data_store_id=data_store_id,\n",
        "   max_documents=30,\n",
        ")\n",
        "results_many = vertex_search_retriever_many.invoke(query)\n",
        "\n",
        "\n",
        "llm_compression = VertexAI(temperature=0., model_name=\"gemini-1.5-flash-001\")\n",
        "chain_filter = LLMChainFilter.from_llm(llm=llm_compression)\n",
        "\n",
        "\n",
        "results_filtered_many = chain_filter.compress_documents(results_many, query)\n",
        "print(len(results_many), len(results_filtered_many))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "qXG5B5-fqr5K",
      "metadata": {
        "id": "qXG5B5-fqr5K"
      },
      "source": [
        "And that's how you adjust the original chain to include filtering into your RAG:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ofmGzdfHhm6t",
      "metadata": {
        "id": "ofmGzdfHhm6t"
      },
      "outputs": [],
      "source": [
        "from langchain_core.runnables import RunnableLambda\n",
        "\n",
        "chain = (\n",
        "   {\"context\": vertex_search_retriever, \"question\": RunnablePassthrough()}\n",
        "   | RunnableLambda(lambda x: {\"context\": chain_filter.compress_documents(x[\"context\"], x[\"question\"]), \"question\": x[\"question\"]})\n",
        "   | prompt\n",
        "   | llm\n",
        "   | StrOutputParser()\n",
        ")\n",
        "\n",
        "\n",
        "chain.invoke(\"How can I make my LLM prompts perform better?\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "utbrGAwpr8Mq",
      "metadata": {
        "id": "utbrGAwpr8Mq"
      },
      "source": [
        "Let's run extraction and compare the length of documents. We can observe a signifant reduction in the overall chunks' (and context) length:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "qwwF7EtCeGQS",
      "metadata": {
        "id": "qwwF7EtCeGQS"
      },
      "outputs": [],
      "source": [
        "from langchain.retrievers.document_compressors import LLMChainExtractor\n",
        "\n",
        "llm_extractor = VertexAI(temperature=0., model_name=\"gemini-1.5-flash-001\")\n",
        "chain_extractor = LLMChainExtractor.from_llm(llm=llm_extractor)\n",
        "\n",
        "\n",
        "results_compressed = chain_filter.compress_documents(result_expansion, query)\n",
        "\n",
        "\n",
        "for original_doc, compressed_doc in zip(result_expansion, results_compressed):\n",
        " print(f\"Document reduced from {len(original_doc.page_content)} to {len(compressed_doc.page_content)}.\")\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
